import cv2 as cv
from ultralytics import YOLO
import os

# 加载模型
model = YOLO('yolov8n.pt')

# 指定图片文件夹路径
image_folder = "D:/python/v8_mine/照片/标准/"
# 指定保存检测到手机的图片文件夹路径
save_folder = "./image_cellphone/"
os.makedirs(save_folder, exist_ok=True)  # 确保保存文件夹存在

# 获取文件夹中所有图片的文件名
image_files = [f for f in os.listdir(image_folder) if f.endswith(('png', 'jpg', 'jpeg'))]

# 遍历图片文件
for image_file in image_files:
    # 构造图片完整路径
    image_path = os.path.join(image_folder, image_file)

    # 读取图片
    img = cv.imread(image_path)

    # 使用YOLO模型进行推理
    results = model(img)

    # 假设results是一个包含单个Results对象的列表（对于单张图片）
    result = results[0]

    # 获取检测结果中的类别ID
    phone_result = model(img)
    phone_result = phone_result[0]
    phone_boxes = phone_result.boxes
    phone_detected = False
    class_ids = phone_boxes.cls.cpu().numpy().astype(int)
    # 检查是否检测到手机（类别ID为67）
    if 67 in class_ids:
        print(f"在图片 {image_file} 中检测到手机")

        # 设置输出路径
        output_path = os.path.join(save_folder, image_file)

        # 保存原始图片到指定文件夹
        cv.imwrite(output_path, img)

        print(f"已保存图片到 {output_path}")

print("所有图片处理完毕")